package com.piece.core.session.repository;

import com.piece.core.framework.support.dto.UserDetailDTO;
import com.piece.core.session.http.RedisSessionExpirationPolicy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationEvent;
import org.springframework.core.NestedExceptionUtils;
import org.springframework.dao.NonTransientDataAccessException;
import org.springframework.data.redis.connection.Message;
import org.springframework.data.redis.connection.MessageListener;
import org.springframework.data.redis.core.BoundHashOperations;
import org.springframework.data.redis.core.RedisOperations;
import org.springframework.expression.Expression;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.session.FindByIndexNameSessionRepository;
import org.springframework.session.MapSession;
import org.springframework.session.Session;
import org.springframework.session.data.redis.RedisFlushMode;
import org.springframework.session.events.SessionCreatedEvent;
import org.springframework.session.events.SessionDeletedEvent;
import org.springframework.session.events.SessionExpiredEvent;
import org.springframework.util.Assert;
import java.time.Duration;
import java.time.Instant;
import java.util.*;
import java.util.Map.Entry;

@Slf4j
public class TokenSessionRepository extends SessionRepository implements FindByIndexNameSessionRepository<TokenSessionRepository.RedisSession>, MessageListener {

    static PrincipalNameResolver PRINCIPAL_NAME_RESOLVER = new PrincipalNameResolver();
    private RedisSessionExpirationPolicy expirationPolicy;

    public TokenSessionRepository() {
        this.redisFlushMode = RedisFlushMode.ON_SAVE;
    }

    public void setSessionRedisOperations(RedisOperations<Object, Object> sessionRedisOperations) {
        Assert.notNull(sessionRedisOperations, "sessionRedisOperations cannot be null");
        this.sessionRedisOperations = sessionRedisOperations;
        this.expirationPolicy = new RedisSessionExpirationPolicy(sessionRedisOperations, this::getExpirationsKey, this::getSessionKey);
        this.configureSessionChannels();
    }

    public void setRedisFlushMode(RedisFlushMode redisFlushMode) {
        Assert.notNull(redisFlushMode, "redisFlushMode cannot be null");
        this.redisFlushMode = redisFlushMode;
    }

    public void setDatabase(int database) {
        this.database = database;
        this.configureSessionChannels();
    }

    private void configureSessionChannels() {
        this.sessionCreatedChannelPrefix = this.redisNamespace + "event:" + this.database + ":created:";
        this.sessionDeletedChannel = "__keyevent@" + this.database + "__:del";
        this.sessionExpiredChannel = "__keyevent@" + this.database + "__:expired";
    }

    public RedisOperations<Object, Object> getSessionRedisOperations() {
        return this.sessionRedisOperations;
    }

    public void save(RedisSession session) {
        session.saveDelta();
        if (session.isNew()) {
            String sessionCreatedKey = this.getSessionCreatedChannel(session.getId());
            this.sessionRedisOperations.convertAndSend(sessionCreatedKey, session.delta);
            session.setNew(false);
        }

    }

    public void cleanupExpiredSessions() {
        this.expirationPolicy.cleanExpiredSessions();
    }

    public RedisSession findById(String id) {
        return this.getSession(id, false);
    }

    public Map<String, RedisSession> findByIndexNameAndIndexValue(String indexName, String indexValue) {
        if (!PRINCIPAL_NAME_INDEX_NAME.equals(indexName)) {
            return Collections.emptyMap();
        } else {
            String principalKey = this.getPrincipalKey(indexValue);
            Set<Object> sessionIds = this.sessionRedisOperations.boundSetOps(principalKey).members();
            Map<String, RedisSession> sessions = new HashMap(sessionIds.size());
            Iterator var6 = sessionIds.iterator();

            while (var6.hasNext()) {
                Object id = var6.next();
                RedisSession session = this.findById((String) id);
                if (session != null) {
                    sessions.put(session.getId(), session);
                }
            }

            return sessions;
        }
    }

    private RedisSession getSession(String id, boolean allowExpired) {
        Map<Object, Object> entries = this.getSessionBoundHashOperations(id).entries();
        if (entries.isEmpty()) {
            return null;
        } else {
            MapSession loaded = this.loadSession(id, entries);
            if (!allowExpired && loaded.isExpired()) {
                return null;
            } else {
                RedisSession result = new RedisSession(loaded);
                result.originalLastAccessTime = loaded.getLastAccessedTime();
                return result;
            }
        }
    }

    private MapSession loadSession(String id, Map<Object, Object> entries) {
        MapSession loaded = new MapSession(id);
        Iterator var4 = entries.entrySet().iterator();

        while (var4.hasNext()) {
            Entry<Object, Object> entry = (Entry) var4.next();
            String key = (String) entry.getKey();
            if ("creationTime".equals(key)) {
                loaded.setCreationTime(Instant.ofEpochMilli((Long) entry.getValue()));
            } else if ("maxInactiveInterval".equals(key)) {
                loaded.setMaxInactiveInterval(Duration.ofSeconds((long) (Integer) entry.getValue()));
            } else if ("lastAccessedTime".equals(key)) {
                loaded.setLastAccessedTime(Instant.ofEpochMilli((Long) entry.getValue()));
            } else if (key.startsWith("sessionAttr:")) {
                loaded.setAttribute(key.substring("sessionAttr:".length()), entry.getValue());
            }
        }

        return loaded;
    }

    public void deleteById(String sessionId) {
        RedisSession session = this.getSession(sessionId, true);
        if (null != session) {
            this.cleanupPrincipalIndex(session);
            this.expirationPolicy.onDelete(session);
            String expireKey = this.getExpiredKey(session.getId());
            this.sessionRedisOperations.delete(expireKey);
            session.setMaxInactiveInterval(Duration.ZERO);
            this.save(session);
        }
    }

    public RedisSession createSession() {
        RedisSession redisSession = new RedisSession();
        if (null != this.maxInactiveIntervalInSeconds) {
            redisSession.setMaxInactiveInterval(Duration.ofSeconds((long) this.maxInactiveIntervalInSeconds));
        }

        return redisSession;
    }

    public RedisSession createSession(String sessionId) {
        RedisSession redisSession = new RedisSession(sessionId);
        if (null != this.maxInactiveIntervalInSeconds) {
            redisSession.setMaxInactiveInterval(Duration.ofSeconds((long) this.maxInactiveIntervalInSeconds));
        }
        return redisSession;
    }

    public RedisSession createSession(String sessionId, long expireTime) {
        RedisSession redisSession = new RedisSession(sessionId);
        redisSession.setMaxInactiveInterval(Duration.ofSeconds(expireTime));
        return redisSession;
    }

    public RedisSession changeSession(String sessionId) {
        RedisSession redisSession = this.getSession(sessionId, false);
        if (null != redisSession) {
            redisSession.changeSessionId(sessionId);
        }
        return redisSession;
    }

    public RedisSession changeSession(String sessionId, long expireTime) {
        RedisSession redisSession = this.getSession(sessionId, false);
        if (null != redisSession) {
            redisSession.setMaxInactiveInterval(Duration.ofSeconds(expireTime));
            redisSession.changeSessionId(sessionId);
        }
        return redisSession;
    }

    public void onMessage(Message message, byte[] pattern) {
        byte[] messageChannel = message.getChannel();
        byte[] messageBody = message.getBody();
        String channel = new String(messageChannel);
        if (channel.startsWith(this.sessionCreatedChannelPrefix)) {
            Map<Object, Object> loaded = (Map) this.defaultRedisSerializer.deserialize(message.getBody());
            this.handleCreated(loaded, channel);
        } else {
            String body = new String(messageBody);
            if (body.startsWith(this.getExpiredKeyPrefix())) {
                boolean isDeleted = channel.equals(this.sessionDeletedChannel);
                if (isDeleted || channel.equals(this.sessionExpiredChannel)) {
                    int beginIndex = body.lastIndexOf(":") + 1;
                    int endIndex = body.length();
                    String sessionId = body.substring(beginIndex, endIndex);
                    RedisSession session = this.getSession(sessionId, true);
                    if (null == session) {
                        log.warn("Unable to publish SessionDestroyedEvent for session " + sessionId);
                        return;
                    }

                    if (log.isDebugEnabled()) {
                        log.debug("Publishing SessionDestroyedEvent for session " + sessionId);
                    }

                    this.cleanupPrincipalIndex(session);
                    if (isDeleted) {
                        this.handleDeleted(session);
                    } else {
                        this.handleExpired(session);
                    }
                }
            }
        }
    }

    private void cleanupPrincipalIndex(RedisSession session) {
        String sessionId = session.getId();
        String principal = PRINCIPAL_NAME_RESOLVER.resolvePrincipal(session, true);
        if (null != principal) {
            this.sessionRedisOperations.boundSetOps(this.getPrincipalKey(principal)).remove(new Object[]{sessionId});
        }
    }

    private void handleCreated(Map<Object, Object> loaded, String channel) {
        String id = channel.substring(channel.lastIndexOf(":") + 1);
        Session session = this.loadSession(id, loaded);
        this.publishEvent(new SessionCreatedEvent(this, session));
    }

    private void handleDeleted(RedisSession session) {
        this.publishEvent(new SessionDeletedEvent(this, session));
    }

    private void handleExpired(RedisSession session) {
        this.publishEvent(new SessionExpiredEvent(this, session));
    }

    private void publishEvent(ApplicationEvent event) {
        try {
            this.applicationEventPublisher.publishEvent(event);
        } catch (Throwable var3) {
            log.error("Error publishing " + event + ".", var3);
        }
    }

    public void setRedisKeyNamespace(String namespace) {
        Assert.hasText(namespace, "namespace cannot be null or empty");
        this.redisNamespace = namespace.trim() + ":";
        this.configureSessionChannels();
    }

    String getSessionKey(String sessionId) {
        return this.redisNamespace + "sessions:" + sessionId;
    }

    String getPrincipalKey(String principalName) {
        return this.redisNamespace + "index:" + FindByIndexNameSessionRepository.PRINCIPAL_NAME_INDEX_NAME + ":" + principalName;
    }

    String getExpirationsKey(long expiration) {
        return this.redisNamespace + "expirations:" + expiration;
    }

    private String getExpiredKey(String sessionId) {
        return this.getExpiredKeyPrefix() + sessionId;
    }

    private String getSessionCreatedChannel(String sessionId) {
        return this.getSessionCreatedChannelPrefix() + sessionId;
    }

    private String getExpiredKeyPrefix() {
        return this.redisNamespace + "sessions:expires:";
    }

    public String getSessionCreatedChannelPrefix() {
        return this.sessionCreatedChannelPrefix;
    }

    public String getSessionDeletedChannel() {
        return this.sessionDeletedChannel;
    }

    public String getSessionExpiredChannel() {
        return this.sessionExpiredChannel;
    }

    private BoundHashOperations<Object, Object, Object> getSessionBoundHashOperations(String sessionId) {
        String key = this.getSessionKey(sessionId);
        return this.sessionRedisOperations.boundHashOps(key);
    }

    static String getSessionAttrNameKey(String attributeName) {
        return "sessionAttr:" + attributeName;
    }

    static class PrincipalNameResolver {
        private SpelExpressionParser parser = new SpelExpressionParser();

        PrincipalNameResolver() {
        }

        public String resolvePrincipal(Session session) {
            return resolvePrincipal(session, false);
        }

        public String resolvePrincipal(Session session, boolean original) {
            String principalName = session.getAttribute(FindByIndexNameSessionRepository.PRINCIPAL_NAME_INDEX_NAME);
            if (null == principalName) {
                Object authentication = session.getAttribute("SPRING_SECURITY_CONTEXT");
                if (null != authentication) {
                    if (original) {
                        if (authentication instanceof SecurityContext) {
                            SecurityContext securityContext = (SecurityContext) authentication;
                            UserDetailDTO user = (UserDetailDTO) securityContext.getAuthentication().getPrincipal();
                            return user.getLoginName();
                        }
                    }
                    Expression expression = this.parser.parseExpression("authentication?.name");
                    return expression.getValue(authentication, String.class);
                } else {
                    return null;
                }
            } else {
                return principalName;
            }
        }
    }

    public final class RedisSession implements Session {
        private final MapSession cached;
        private Instant originalLastAccessTime;
        private Map<String, Object> delta;
        private boolean isNew;
        private String originalPrincipalName;
        private String originalSessionId;

        RedisSession() {
            this(new MapSession());
            this.delta.put("creationTime", this.getCreationTime().toEpochMilli());
            this.delta.put("maxInactiveInterval", (int) this.getMaxInactiveInterval().getSeconds());
            this.delta.put("lastAccessedTime", this.getLastAccessedTime().toEpochMilli());
            this.isNew = true;
            this.flushImmediateIfNecessary();
        }

        RedisSession(String sessionId) {
            this(new MapSession(sessionId));
            this.delta.put("creationTime", this.getCreationTime().toEpochMilli());
            this.delta.put("maxInactiveInterval", (int) this.getMaxInactiveInterval().getSeconds());
            this.delta.put("lastAccessedTime", this.getLastAccessedTime().toEpochMilli());
            this.isNew = true;
            this.flushImmediateIfNecessary();
        }

        RedisSession(MapSession cached) {
            this.delta = new HashMap();
            Assert.notNull(cached, "MapSession cannot be null");
            this.cached = cached;
            this.originalPrincipalName = TokenSessionRepository.PRINCIPAL_NAME_RESOLVER.resolvePrincipal(this);
            this.originalSessionId = cached.getId();
        }

        public void setNew(boolean isNew) {
            this.isNew = isNew;
        }

        public void setLastAccessedTime(Instant lastAccessedTime) {
            this.cached.setLastAccessedTime(lastAccessedTime);
            this.putAndFlush("lastAccessedTime", this.getLastAccessedTime().toEpochMilli());
        }

        public boolean isExpired() {
            return this.cached.isExpired();
        }

        public boolean isNew() {
            return this.isNew;
        }

        public Instant getCreationTime() {
            return this.cached.getCreationTime();
        }

        public String getId() {
            return this.cached.getId();
        }

        public String changeSessionId() {
            return this.cached.changeSessionId();
        }

        public void changeSessionId(String sessionId) {
            this.cached.setId(sessionId);
        }

        public Instant getLastAccessedTime() {
            return this.cached.getLastAccessedTime();
        }

        public void setMaxInactiveInterval(Duration interval) {
            this.cached.setMaxInactiveInterval(interval);
            this.putAndFlush("maxInactiveInterval", (int) this.getMaxInactiveInterval().getSeconds());
        }

        public Duration getMaxInactiveInterval() {
            return this.cached.getMaxInactiveInterval();
        }

        public <T> T getAttribute(String attributeName) {
            return this.cached.getAttribute(attributeName);
        }

        public Set<String> getAttributeNames() {
            return this.cached.getAttributeNames();
        }

        public void setAttribute(String attributeName, Object attributeValue) {
            this.cached.setAttribute(attributeName, attributeValue);
            this.putAndFlush(TokenSessionRepository.getSessionAttrNameKey(attributeName), attributeValue);
        }

        public void removeAttribute(String attributeName) {
            this.cached.removeAttribute(attributeName);
            this.putAndFlush(TokenSessionRepository.getSessionAttrNameKey(attributeName), (Object) null);
        }

        private void flushImmediateIfNecessary() {
            if (TokenSessionRepository.this.redisFlushMode == RedisFlushMode.IMMEDIATE) {
                this.saveDelta();
            }
        }

        private void putAndFlush(String a, Object v) {
            this.delta.put(a, v);
            this.flushImmediateIfNecessary();
        }

        private void saveDelta() {
            String sessionId = this.getId();
            this.saveChangeSessionId(sessionId);
            if (!this.delta.isEmpty()) {
                TokenSessionRepository.this.getSessionBoundHashOperations(sessionId).putAll(this.delta);
                String principalSessionKey = TokenSessionRepository.getSessionAttrNameKey(FindByIndexNameSessionRepository.PRINCIPAL_NAME_INDEX_NAME);
                String securityPrincipalSessionKey = TokenSessionRepository.getSessionAttrNameKey("SPRING_SECURITY_CONTEXT");
                if (this.delta.containsKey(principalSessionKey) || this.delta.containsKey(securityPrincipalSessionKey)) {
                    String principal;
                    if (this.originalPrincipalName != null) {
                        principal = TokenSessionRepository.this.getPrincipalKey(this.originalPrincipalName);
                        TokenSessionRepository.this.sessionRedisOperations.boundSetOps(principal).remove(new Object[]{sessionId});
                    }

                    principal = TokenSessionRepository.PRINCIPAL_NAME_RESOLVER.resolvePrincipal(this, true);
                    this.originalPrincipalName = principal;
                    if (principal != null) {
                        String principalRedisKey = TokenSessionRepository.this.getPrincipalKey(principal);
                        TokenSessionRepository.this.sessionRedisOperations.boundSetOps(principalRedisKey).add(new Object[]{sessionId});
                    }
                }

                this.delta = new HashMap(this.delta.size());
                if (this.originalLastAccessTime == null) {
                    this.originalLastAccessTime = Instant.now();
                }
                Long originalExpiration = this.originalLastAccessTime.plus(this.getMaxInactiveInterval()).toEpochMilli();
                TokenSessionRepository.this.expirationPolicy.onExpirationUpdated(originalExpiration, this);
            }
        }

        private void saveChangeSessionId(String sessionId) {
            if (!sessionId.equals(this.originalSessionId)) {
                if (!this.isNew()) {
                    String originalSessionIdKey = TokenSessionRepository.this.getSessionKey(this.originalSessionId);
                    String sessionIdKey = TokenSessionRepository.this.getSessionKey(sessionId);
                    TokenSessionRepository.this.sessionRedisOperations.rename(originalSessionIdKey, sessionIdKey);
                    String originalExpiredKey = TokenSessionRepository.this.getExpiredKey(this.originalSessionId);
                    String expiredKey = TokenSessionRepository.this.getExpiredKey(sessionId);

                    try {
                        TokenSessionRepository.this.sessionRedisOperations.rename(originalExpiredKey, expiredKey);
                    } catch (NonTransientDataAccessException var7) {
                        if (!"ERR no such key".equals(NestedExceptionUtils.getMostSpecificCause(var7).getMessage())) {
                            throw var7;
                        }
                    }
                }

                this.originalSessionId = sessionId;
            }
        }
    }
}
